import sys
import uuid
import time
import pickle
import os

import jax
import jax.numpy as jnp
import brax
import optax
from robotisgp_gradient import RobotisgpGradient
from brax.envs.fetch import Fetch
from brax.io import html

def generate_traj_loss(s1, n, env_step):
    def traj_loss(a_list):
        reward = 0
        st = s1
        for i in range(n):
            st = env_step(st, a_list[i])
            reward += st.reward
        return -reward
    return traj_loss

def get_distance(st, s0, mass):
    pos_before = s0.qp.pos[:-1]  # ignore floor at last index
    pos_after = st.qp.pos[:-1]  # ignore floor at last index
    com_before = jnp.sum(pos_before * mass, axis=0) / jnp.sum(mass)
    com_after = jnp.sum(pos_after * mass, axis=0) / jnp.sum(mass)
    d = com_after[1] - com_before[1]
    return d

T = 500
rng = jax.random.PRNGKey(1)
key, rng = jax.random.split(rng)
logdir = sys.argv[1]
env_fn = Fetch
env = env_fn()
s0 = env.reset(key)
st = s0
jit_step_fn = jax.jit(env.step)

if os.path.exists(f"{logdir}/a1_n.pickle"):
    with open(f"{logdir}/a1_n.pickle", "rb") as f:
        a1_n = pickle.load(f)
else:
    a1_n = dict()
    for i in range(T):
        key, rng = jax.random.split(rng)
        a1_n[i] = jax.random.normal(key, (env.action_size, ))
print("initialized a1_n")

loss = generate_traj_loss(s0, T, jit_step_fn)

def fit(a1_n, optimizer):
    opt_state = optimizer.init(a1_n)

    def step(a1_n, opt_state):
        loss_value, grads = jax.value_and_grad(loss)(a1_n)
        updates, opt_state = optimizer.update(grads, opt_state, a1_n)
        a1_n = optax.apply_updates(a1_n, updates)
        return a1_n, opt_state, loss_value

    for i in range(240):
        a1_n, opt_state, loss_value = step(a1_n, opt_state)
        print(f'step {i}, loss: {loss_value}')

    return a1_n

optimizer = optax.adam(learning_rate=1e-1)
a1_n = fit(a1_n, optimizer)
with open(f"{logdir}/a1_n.pickle","wb") as f:
    pickle.dump(a1_n, f)

t = 0
qps = []
while t<T:
    qps.append(st.qp)
    act = a1_n[t]
    st = jit_step_fn(st, act)
    t += 1
#print(get_distance(st, s0, env.mass))
#print(st.qp.pos[0, 2])

html_path = f'{logdir}/trajectory_{uuid.uuid4()}.html'
html.save_html(html_path, env.sys, qps, True)
